import torch
from RGCNmodel import RGCN
from config import *
from sklearn.metrics import roc_auc_score
from datetime import datetime
from tqdm import tqdm
from RGCNDataset import RGCNDataset
from dataPre import dataPre
from torch.utils.data import DataLoader
import sys
sys.path.append(r"/home/cyw/projects/function_sim_project/basic_script")


def write_to_logs(inf):
    with open("/home/cyw/projects/function_sim_project/R_GCN/RGCN_logs.txt", 'a+') as file:
        data = str(datetime.now())+"\t"+inf+"\n"
        file.write(data)
        print(data)


class RGCNModelWarpper():
    def __init__(self, model=None) -> None:
        # seed_everything()
        if model == None:
            self.model = RGCN(
                embSize=embSize, wordSize=wordsize, layerSize=laySize)
            write_to_logs("模型参数：lr:{}    epoch:{}    laysize:{}    embsize:{}    margin:{}\n".format(
                lr, epoch, laySize, embSize, margin))
        else:
            self.model = model

        self.optimizer = torch.optim.Adam(
            [{'params': self.model.parameters()},], lr=lr)

    def get_model_dataloader(self, datas, name, is_train, is_shuffle=True):
        """
            else_datas用于加载其它的输入,方便其他模型的对比
        """
        my_dataset = RGCNDataset(datas, name, is_train=is_train)
        dataloader = DataLoader(my_dataset, batch_size=batchSize, shuffle=is_shuffle, num_workers=40,
                                collate_fn=my_dataset.adjust_samples_to_same_dimension)
        return dataloader

    def train(self, dataloader, valid_dataloader, validValues):
        print("----------------R-GCN开始训练----------------")
        totalLoss = "-"
        best_auc = -float("inf")
        for e in tqdm(range(epoch)):
            self.model.train()
            temploss = 0
            for i, (query, target1, target2, target3) in enumerate(dataloader):
                torch.cuda.empty_cache()
                self.optimizer.zero_grad()

                adj_x, att_x, vtype_x = target1
                adj_y, att_y, vtype_y = target2
                adj_z, att_z, vtype_z = target3

                adj_x, adj_y, adj_z = adj_x.to(
                    device), adj_y.to(device), adj_z.to(device)
                att_x, att_y, att_z = att_x.to(
                    device), att_y.to(device), att_z.to(device)
                vtype_x, vtype_y, vtype_z = vtype_x.to(
                    device), vtype_y.to(device), vtype_z.to(device)

                score1 = self.model(adj_x, att_x, vtype_x,
                                    adj_y, att_y, vtype_y)
                score2 = self.model(adj_x, att_x, vtype_x,
                                    adj_z, att_z, vtype_z)

                loss = torch.relu(score2 - score1 + margin)
                loss = torch.mean(loss)

                temploss += loss
                loss.backward()
                self.optimizer.step()
            temp_auc = self.test_model_auc(valid_dataloader, validValues)
            write_to_logs("epoch:{} auc:{} loss:{}".format(
                e, temp_auc, temploss))
            if temp_auc > best_auc:
                best_auc = temp_auc
                torch.save(
                    self.model, r'/home/cyw/projects/function_sim_project/all_data/models/R_GCN_model')
                write_to_logs("模型已保存")
            totalLoss += str(float(temploss))
            totalLoss += "-"
        print("训练完成")

    def test_model_auc(self, dataloader, values):
        ans = self.get_model_score(dataloader)
        auc_score = roc_auc_score(values, ans)
        return auc_score

    def get_model_score(self, dataloader):
        """
            输入指定格式数据，通过训练好的模型得到相应的得分
            把dataloader输入进来，方便之前做数据分布的处理
            不为train了，则数据只有前三个是有用的，第四个出现值说明代码错误
        """
        self.model.eval()
        with torch.no_grad():
            res = []
            for i, (query, target1, target2, target3) in enumerate(dataloader):
                torch.cuda.empty_cache()
                self.optimizer.zero_grad()
                adj_x, att_x, vtype_x = target1
                adj_y, att_y, vtype_y = target2
                adj_x = adj_x.to(device)
                adj_y = adj_y.to(device)
                att_x = att_x.to(device)
                att_y = att_y.to(device)
                vtype_x = vtype_x.to(device)
                vtype_y = vtype_y.to(device)
                score1 = self.model(adj_x, att_x, vtype_x,
                                    adj_y, att_y, vtype_y)  # B*1
                res.append(score1)
            ans = torch.cat((res))
        return ans.tolist()


if __name__ == "__main__":
    datapre = dataPre()
    data_pair_infs = datapre.get_pair_infs()

    testvalues, validvalues = [], []
    for i in range(len(data_pair_infs["test"])):
        testvalues.append(data_pair_infs["test"][i][2])
    for i in range(len(data_pair_infs["valid"])):
        validvalues.append(data_pair_infs["valid"][i][2])

    print("RGCN正在训练模型")
    RGCNTrainer = RGCNModelWarpper()
    test_dataloader = RGCNTrainer.get_model_dataloader(
        data_pair_infs["test"], "test", is_train=False, is_shuffle=False)
    train_dataloader = RGCNTrainer.get_model_dataloader(
        data_pair_infs["train"], "train", is_train=True)
    valid_dataloader = RGCNTrainer.get_model_dataloader(
        data_pair_infs["valid"], "valid", is_train=False, is_shuffle=False)
    RGCNTrainer.train(train_dataloader, valid_dataloader, validvalues)
